from benchmark.config import LocalCommittee, NodeParameters, BenchParameters, ConfigError
from benchmark.utils import BenchError, Print, PathMaker
from benchmark.commands import CommandMaker
from benchmark.logs import LogParser, ParseError
import subprocess
import time

class LocalBench:
    def __init__(self, bench_parameters_dict, node_parameters_dict):
        try:
            self.bench_parameters = BenchParameters(bench_parameters_dict)
            self.node_parameters = NodeParameters(node_parameters_dict)
        except ConfigError as e:
            raise BenchError('Invalid nodes or bench parameters', e)

    def run(self, debug=False):
        Print.heading('Starting local benchmark')
        
        # Cleanup
        cmd = CommandMaker.cleanup_configs()
        subprocess.run([cmd], shell=True, stderr=subprocess.DEVNULL)
        
        # Compile
        cmd = CommandMaker.compile().split()
        subprocess.run(cmd, check=True)
        
        # Generate keys
        nodes = self.bench_parameters.nodes[0]
        key_files = [PathMaker.key_file(i) for i in range(nodes)]
        cmd = CommandMaker.generate_key(path="./", nodes=nodes).split()
        subprocess.run(cmd, check=True)
        
        # Generate committee
        keys = []
        for filename in key_files:
            with open(filename, 'r') as f:
                import json
                data = json.load(f)
                keys.append(data['public'])
        
        committee = LocalCommittee(keys, list(range(nodes)), 8000)
        committee.print(PathMaker.committee_file())
        
        # Generate parameters
        self.node_parameters.print(PathMaker.parameters_file())
        
        # Run nodes
        dbs = [PathMaker.db_path(i) for i in range(nodes)]
        node_logs = [PathMaker.node_log_info_file(i, "local") for i in range(nodes)]
        
        for i in range(nodes):
            cmd = CommandMaker.run_node(
                i,
                key_files[i],
                PathMaker.threshold_key_file(i),
                PathMaker.committee_file(),
                dbs[i],
                PathMaker.parameters_file(),
                "local",
                self.bench_parameters.log_level
            )
            # Run in background
            subprocess.Popen(cmd, shell=True)
        
        # Wait for benchmark
        Print.info(f'Running benchmark for {self.bench_parameters.duration} seconds...')
        time.sleep(self.bench_parameters.duration)
        
        # Kill processes
        subprocess.run([CommandMaker.kill()], shell=True, stderr=subprocess.DEVNULL)
        
        # Parse results
        Print.info('Parsing logs...')
        try:
            parser = LogParser.process('./logs/local', 
                                    faults=self.node_parameters.faults,
                                    protocol=self.bench_parameters.protocol,
                                    ddos=self.node_parameters.ddos)
            return parser
        except ParseError as e:
            raise BenchError('Failed to parse logs', e)
